# -*- coding: utf-8 -*-
#source https://github.com/SunZhimin2021/AIPentest/blob/main/Tools/autoexploit.py
#Author Sunzhimin
#公众号 AI与安全

import nmap
import json
from pymetasploit3.msfrpc import MsfRpcClient
import subprocess
import time
import argparse
import sys
from openai import OpenAI  # deepseek也用openai接口
from dotenv import load_dotenv

load_dotenv()
api_key = os.getenv("api_key")
base_url=os.getenv('base_url')
modelname=os.getenv('modelname')

class AutoExploiter:
    # 预定义prompt模板
    ANALYSIS_PROMPT_TEMPLATE = """分析以下Nmap扫描结果，并提供：
1. 可能存在的漏洞
2. 推荐使用的Metasploit模块（包括完整路径）
3. 必要的参数设置
4. 利用的成功概率评估

扫描结果：
{scan_results}

请以JSON格式返回，格式如下：
{{
    "exploits": [
        {{
            "type": "exploit/auxiliary",
            "path": "完整msf模块路径",
            "port": port_number,
            "params": {{"参数名": "参数值"}},
            "probability": "成功概率评估",
            "description": "漏洞描述"
            "msfconsolecommand"："可以在命令行直接执行的msfconsole命令，格式为 msfconsole -x ...."
        }}
    ]
}}

严格按照上述JSON格式返回，不要包含任何其他文字说明。也不需要```json符号，以便解析。
"""

    def __init__(self, target_ip, lhost):
        self.target_ip = target_ip
        self.lhost = lhost
        self.scan_results = None
        self.msf_client = None
        self.deepseek_client = OpenAI(api_key=api_key,base_url=base_url)  # 使用OpenAI客户端

    
    def analyze_with_gpt(self):
        """使用GPT分析扫描结果"""
        try:
            # 格式化扫描结果
            formatted_results = json.dumps(self.scan_results, indent=2)
            
            # 使用预定义的prompt模板，填入扫描结果
            prompt = self.ANALYSIS_PROMPT_TEMPLATE.format(
                scan_results=formatted_results
            )
            
            # 调用OpenAI API
            response = self.deepseek_client.chat.completions.create(
                #model="deepseek-v3",  # 阿里云用的是模型名称，deepseek官网用的是deepseek-chat
                model=modelname,
                messages=[
                    {"role": "system", "content": "You are a cybersecurity expert specialized in vulnerability analysis and exploitation."},
                    {"role": "user", "content": prompt}
                ],
                temperature=0
            )
            print(response.choices[0].message.content)
            # 解析响应
            try:
                analysis = json.loads(response.choices[0].message.content)
                print("\n[+] GPT分析完成")
                return analysis.get('exploits', [])
            except json.JSONDecodeError:
                print("[-] 无法解析GPT的响应")
                print("响应内容:", response.choices[0].message.content)
                return []
                
        except Exception as e:
            print(f"[-] GPT分析失败: {str(e)}")
            return []

    def scan_target(self):
        """使用Nmap扫描目标"""
        try:
            print(f"[*] 开始扫描目标 {self.target_ip}...")
            
            nm = nmap.PortScanner()
            #nm.scan(self.target_ip, arguments='-sV -sC -O -p- --script vuln')
            nm.scan(self.target_ip, arguments='-sV ')
            self.scan_results = nm[self.target_ip]
            print("[+] 扫描完成")
            print(self.scan_results)
            return self.analyze_with_gpt()  # 改为调用GPT分析
        except Exception as e:
            print(f"[-] 扫描失败: {str(e)}")
            return None
    
# main()函数中的修改
def main():
    parser = argparse.ArgumentParser(description='自动化漏洞扫描与利用工具')
    parser.add_argument('-t', '--target', required=True, help='目标IP地址')
    parser.add_argument('-l', '--lhost', required=True, help='本地IP地址(用于接收反弹shell)')
    
    args = parser.parse_args()

    exploiter = AutoExploiter(args.target, args.lhost)
    
    try:
        
        
        # 扫描目标并使用GPT分析结果
        print('scan beginning')
        exploits = exploiter.scan_target()
        
        if not exploits:
            print("[-] 未发现可利用的漏洞")
            return
        
        print("\n[+] GPT分析发现以下可能的漏洞利用方法:")
        for i, exploit in enumerate(exploits):
            print(f"{i+1}. {exploit['path']}")
            print(f"   描述: {exploit['description']}")
            print(f"   端口: {exploit['port']}")
            print(f"   成功率: {exploit['probability']}")
            print(f"   msf命令: {exploit['msfconsolecommand']}")
            print()
        
        # 让用户选择要使用的漏洞利用
       
    finally:
        print('finish')
        
if __name__ == "__main__":
    main()
